import tensorflow as tf
from tensorflow.keras import layers, optimizers, Sequential, losses
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
from dataset_self import *
from specNet import SpecVAE
import colour


Set_GPU_Memory_Growth()  # GPU按需分配

log_dir = './log_AE/AE_demo/'
sum_w = tf.summary.create_file_writer(log_dir)
path = '.\\result_duplicate\\all_para'
x_train, y_train, x_test, y_test, lam = res_to_data(get_R(path), 0.2)  # 生成数据集， 将json文件中的dict转为TensorSlice类型数据

# pos_train, pos_test = get_xyz_pos([y_train, y_test])  # 计算生成对应反射谱的颜色坐标
pos_train, pos_test = np.load('pos_train.npy'), np.load('pos_test.npy')
norm = NormPara(x_train, pos_train)  # 归一化预处理
train_da = tf.data.Dataset.from_tensor_slices((x_train, pos_train))
test_da = tf.data.Dataset.from_tensor_slices((x_test, pos_test))
train_db = train_da.map(norm.preprocess).shuffle(10000).batch(3000)  # 训练数据集，打乱，批量
test_db = test_da.map(norm.process).shuffle(10000).batch(1000)  # 测试数据集，打乱，批量

net = SpecVAE()

epochs = 2000
lr = tf.keras.optimizers.schedules.PiecewiseConstantDecay([500, 1000], [5e-5, 5e-6, 1e-6])
# lr = 1e-5
opti = optimizers.Adam(learning_rate=lr)

# 开始训练
loss_array = []
ep_array = []
ep_array2 = []
test_array = []
loss_test = 0
for ep in range(epochs):
    loss = 0
    for i, (x1, x2) in enumerate(train_db):
        with tf.GradientTape() as tape:
            y_hat, mu, sig = net(x1, x2)
            loss_out = tf.reduce_mean(losses.mean_squared_error(x1, y_hat))
            loss_mid = tf.reduce_mean(-0.5 * (sig + 1 - mu ** 2 - tf.exp(sig)))  # 中间的KL散度
            loss += loss_out + loss_mid
            var = net.trainable_variables
        grads = tape.gradient(loss, var)
        opti.apply_gradients(zip(grads, var))
    loss = loss / x1.shape[0]
    loss_array.append(loss)
    ep_array.append(ep)

    # 测试集表现评估
    if ep % 50 == 1:
        loss_test = 0
        for j, (t1, t2) in enumerate(test_db):
            y_p, mu2, sig2 = net(t1, t2)
            loss_out_test = tf.reduce_mean(losses.mean_squared_error(t1, y_p))
            loss_mid_test = tf.reduce_mean(-0.5 * (sig2 + 1 - mu2 ** 2 - tf.exp(sig2)))
            loss_test += loss_out_test + loss_mid_test
        loss_test = loss_test / t1.shape[0]
        test_array.append(loss_test)
        ep_array2.append(ep)

    print('\r' + 'epoch: %d/%d  train loss: %.5f test loss: %.5f' % (ep, epochs, loss, loss_test), end='')
